package ch.trainingtest.utils;

import java.io.InputStream;
import java.sql.Connection;
import java.sql.SQLException;

import javax.persistence.EntityManager;

import org.apache.commons.lang.StringUtils;
import org.dbunit.DatabaseUnitException;
import org.dbunit.database.DatabaseConfig;
import org.dbunit.database.DatabaseConnection;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.dataset.DataSetException;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSetBuilder;
import org.dbunit.operation.DatabaseOperation;
public class DbUnitHelper {
	
    private IDatabaseConnection connection;
    
    public DbUnitHelper(EntityManager entityManager, String schema) throws DatabaseUnitException  {
    		setConnection(entityManager, schema);	
    }  

	public final void setConnection(EntityManager entityManager, String schema) throws DatabaseUnitException {	
		entityManager.getTransaction().begin();
		Connection jdbcConnection = entityManager.unwrap(Connection.class);
		entityManager.getTransaction().commit();
		if (!StringUtils.trimToEmpty(schema).equals("")){
			connection = new DatabaseConnection(jdbcConnection, schema, true);
		} else {
			connection = new DatabaseConnection(jdbcConnection);
		}
	    connection.getConfig().setProperty(DatabaseConfig.FEATURE_QUALIFIED_TABLE_NAMES, false);		
	}	 
    
    /**
     * get dataSet from XML file
     * @param fileName
     * @return
     * @throws Exception
     */
    private IDataSet getDataSet(String fileName) throws DataSetException {
    	InputStream flatFile = getClass().getResourceAsStream(fileName);
    	IDataSet dataSet = new FlatXmlDataSetBuilder().build(flatFile);
    	return dataSet;
    }
    
    /**
     * Loads the DBUnit test data.
     * @param fileName  The name of the XML DBUnit input file
	 * @throws Exception 
     */
    private void insertData(IDataSet dataSet, boolean isClearData) throws Exception {
        if (isClearData) {
        	DatabaseOperation.CLEAN_INSERT.execute(connection, dataSet);
        } else {
        	DatabaseOperation.INSERT.execute(connection, dataSet);
        }
    }
    
    /**
     * Loads the DBUnit test data.
     * @param fileName  The name of the XML DBUnit input file
	 * @throws Exception 
     */
    public void loadTestData(String fileName, boolean isClearData) throws Exception {
        IDataSet dataSet = getDataSet(fileName);
        insertData(dataSet, isClearData) ;         
    }
    
    /**
     * Excuse clear data.
     */
    public void cleanData(String fileName) throws DataSetException, DatabaseUnitException, SQLException  {    	
    	IDataSet dataSet = getDataSet(fileName);
    	DatabaseOperation.DELETE.execute(connection, dataSet);
    }
    
}